# -*- coding: utf-8 -*-
"""
数据库模块 - 定义表结构和基础操作
"""

import sqlite3
import os
from config import DB_PATH, DEFAULT_DEVICE_TYPES, DEFAULT_NETWORK_SEGMENTS, DEFAULT_MIDDLEWARE_TYPES, DEFAULT_MIDDLEWARE_NAMES, DEFAULT_VIRTUALIZATION_TYPES, DEFAULT_ACCOUNT_LEVELS, DEFAULT_ACCOUNT_TYPES, DEFAULT_APP_SYSTEM_TYPES
from utils import ensure_dir

class Database:
    def __init__(self):
        self.db_path = DB_PATH
        self.init_database()
    
    def get_connection(self):
        """获取数据库连接"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn
    
    def init_database(self):
        """初始化数据库表结构"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        # 创建IP管理表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS ip_management (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                ip_address TEXT UNIQUE NOT NULL,
                device_type TEXT NOT NULL,
                department TEXT NOT NULL,
                purpose TEXT NOT NULL,
                asset_id TEXT,
                responsible_person TEXT,
                network_segment TEXT,
                status TEXT DEFAULT '启用',  -- 状态：启用/禁用
                remark TEXT,  -- 备注
                created_time TEXT DEFAULT CURRENT_TIMESTAMP,
                updated_time TEXT DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 检查并添加新的列（如果它们不存在）
        try:
            # 检查是否已存在status和remark列
            cursor.execute("SELECT * FROM ip_management LIMIT 0")
            columns = [description[0] for description in cursor.description]
            
            if 'status' not in columns:
                cursor.execute("ALTER TABLE ip_management ADD COLUMN status TEXT DEFAULT '启用'")
            
            if 'remark' not in columns:
                cursor.execute("ALTER TABLE ip_management ADD COLUMN remark TEXT")
                
        except Exception:
            # 如果表是空的，可能无法获取列信息，这种情况下我们假设列已存在
            pass
        
        # 创建资产信息表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS asset_info (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                asset_number TEXT UNIQUE NOT NULL,
                device_type TEXT NOT NULL,
                model TEXT,
                purchase_date TEXT,
                cpu_cores INTEGER,
                memory_gb REAL,
                disk_gb REAL,
                virtualization_type TEXT,
                created_time TEXT DEFAULT CURRENT_TIMESTAMP,
                updated_time TEXT DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建中间件管理表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS middleware_info (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                type TEXT NOT NULL,
                version TEXT,
                install_directory TEXT,
                port INTEGER,
                associated_ip TEXT,
                associated_system TEXT,
                created_time TEXT DEFAULT CURRENT_TIMESTAMP,
                updated_time TEXT DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建应用系统管理表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS app_system (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                system_name TEXT NOT NULL,
                login_url TEXT,
                admin_account TEXT,
                admin_password TEXT,
                associated_ip TEXT,
                system_type TEXT,
                created_time TEXT DEFAULT CURRENT_TIMESTAMP,
                updated_time TEXT DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建操作系统账号管理表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS os_accounts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                account_name TEXT NOT NULL,
                account_level TEXT NOT NULL,
                account_type TEXT NOT NULL,
                account_password TEXT,
                associated_ip TEXT,
                created_time TEXT DEFAULT CURRENT_TIMESTAMP,
                updated_time TEXT DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建字典表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS dictionary (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                dict_name TEXT NOT NULL,
                dict_value TEXT NOT NULL,
                description TEXT,
                created_time TEXT DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 创建审计日志表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS audit_log (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                operation_time TEXT DEFAULT CURRENT_TIMESTAMP,
                operator TEXT DEFAULT 'admin',
                operation_type TEXT NOT NULL,
                operation_object TEXT,
                operation_detail TEXT
            )
        ''')
        
        # 初始化字典数据（仅在字典表为空时）
        cursor.execute('SELECT COUNT(*) FROM dictionary')
        dict_count = cursor.fetchone()[0]
        if dict_count == 0:
            self.init_dictionary_data(cursor)
        
        conn.commit()
        conn.close()
    
    def init_dictionary_data(self, cursor):
        """初始化字典默认数据"""
        # 检查是否已存在字典数据，避免重复插入
        cursor.execute('SELECT COUNT(*) FROM dictionary')
        dict_count = cursor.fetchone()[0]
        if dict_count > 0:
            return
        
        # 设备类型
        for device_type in DEFAULT_DEVICE_TYPES:
            cursor.execute('''
                INSERT INTO dictionary (dict_name, dict_value, description)
                VALUES (?, ?, ?)
            ''', ('设备类型', device_type, '设备类型字典'))
        
        # 网段
        for segment in DEFAULT_NETWORK_SEGMENTS:
            cursor.execute('''
                INSERT INTO dictionary (dict_name, dict_value, description)
                VALUES (?, ?, ?)
            ''', ('网段', segment, '网络网段字典'))
        
        # 中间件类型
        for middleware_type in DEFAULT_MIDDLEWARE_TYPES:
            cursor.execute('''
                INSERT INTO dictionary (dict_name, dict_value, description)
                VALUES (?, ?, ?)
            ''', ('中间件类型', middleware_type, '中间件类型字典'))
        
        # 中间件名称
        for middleware_name in DEFAULT_MIDDLEWARE_NAMES:
            cursor.execute('''
                INSERT INTO dictionary (dict_name, dict_value, description)
                VALUES (?, ?, ?)
            ''', ('中间件名称', middleware_name, '中间件名称字典'))
        
        # 虚拟化类型
        for virt_type in DEFAULT_VIRTUALIZATION_TYPES:
            cursor.execute('''
                INSERT INTO dictionary (dict_name, dict_value, description)
                VALUES (?, ?, ?)
            ''', ('虚拟化类型', virt_type, '虚拟化类型字典'))
        
        # 账号级别
        for level in DEFAULT_ACCOUNT_LEVELS:
            cursor.execute('''
                INSERT INTO dictionary (dict_name, dict_value, description)
                VALUES (?, ?, ?)
            ''', ('账号级别', level, '账号级别字典'))
        
        # 账号类型
        for acc_type in DEFAULT_ACCOUNT_TYPES:
            cursor.execute('''
                INSERT INTO dictionary (dict_name, dict_value, description)
                VALUES (?, ?, ?)
            ''', ('账号类型', acc_type, '账号类型字典'))
        
        # 应用系统类型
        for app_type in DEFAULT_APP_SYSTEM_TYPES:
            cursor.execute('''
                INSERT INTO dictionary (dict_name, dict_value, description)
                VALUES (?, ?, ?)
            ''', ('应用系统类型', app_type, '应用系统类型字典'))
    
    def log_operation(self, operation_type, operation_object=None, operation_detail=None):
        """记录操作日志"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute('''
            INSERT INTO audit_log (operation_type, operation_object, operation_detail)
            VALUES (?, ?, ?)
        ''', (operation_type, operation_object, operation_detail))
        conn.commit()
        conn.close()
    
    def get_dictionary_values(self, dict_name):
        """获取字典值列表"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute('''
            SELECT DISTINCT dict_value FROM dictionary WHERE dict_name = ? ORDER BY dict_value
        ''', (dict_name,))
        values = [row[0] for row in cursor.fetchall()]
        conn.close()
        return values
    
    def add_dictionary_value(self, dict_name, dict_value, description=None):
        """添加字典值"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute('''
            INSERT INTO dictionary (dict_name, dict_value, description)
            VALUES (?, ?, ?)
        ''', (dict_name, dict_value, description))
        conn.commit()
        conn.close()
        self.log_operation('添加字典值', dict_name, f'{dict_name}: {dict_value}')
    
    def check_ip_exists(self, ip_address, exclude_id=None):
        """检查IP地址是否已存在"""
        conn = self.get_connection()
        cursor = conn.cursor()
        if exclude_id:
            cursor.execute('''
                SELECT COUNT(*) FROM ip_management WHERE ip_address = ? AND id != ?
            ''', (ip_address, exclude_id))
        else:
            cursor.execute('''
                SELECT COUNT(*) FROM ip_management WHERE ip_address = ?
            ''', (ip_address,))
        count = cursor.fetchone()[0]
        conn.close()
        return count > 0